package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
)

type AesCipher struct {
	block cipher.Block
}

func NewAesCipher(key string) (c *AesCipher, err error) {
	block, err := aes.NewCipher([]byte(key))
	if err == nil {
		c = &AesCipher{
			block: block,
		}
	}
	return
}

// See alternate IV creation from ciphertext below
//var iv = []byte{35, 46, 57, 24, 85, 35, 24, 74, 87, 35, 88, 98, 66, 32, 14, 05}

func (a *AesCipher) Encrypt(data []byte) ([]byte, error) {
	ciphertext := make([]byte, aes.BlockSize+len(data))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}
	cfb := cipher.NewCFBEncrypter(a.block, iv)
	cfb.XORKeyStream(ciphertext[aes.BlockSize:], data)
	return ciphertext, nil
}

func (a *AesCipher) EncryptToBase64(data []byte) (rs string, err error) {
	ciphertext, err := a.Encrypt(data)
	if err == nil {
		rs = base64.StdEncoding.EncodeToString(ciphertext)
	}
	return rs, err
}

func (a *AesCipher) Decrypt(text []byte) ([]byte, error) {
	if len(text) < aes.BlockSize {
		return nil, errors.New("ciphertext too short")
	}
	iv := text[:aes.BlockSize]
	text = text[aes.BlockSize:]
	cfb := cipher.NewCFBDecrypter(a.block, iv)
	cfb.XORKeyStream(text, text)
	return text, nil
}

func (a *AesCipher) DecryptBase64(text string) ([]byte, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(text)
	if err != nil {
		return nil, err
	}
	return a.Decrypt(ciphertext)
}
